#include "crypt.h"
#include "rsa_crypt.h"
#include "aes_crypt.h"

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include <error.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <unistd.h>
#include <errno.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <poll.h>

#define SERVER_ADDR "127.0.0.1"
#define SERVER_PORT 5050

#define AES_KEY_FIXED "01234567890123456789012345678901"

static int sock_fd;
static unsigned char aes_key[AES_KEY_LEN];
static RSA *rsa_key;

#define RECV_BUF_SIZE 2048

int crypt_client(void);
int frame_parse(const unsigned char *cmd,int len);

int main(int argc,char *argv[])
{
    crypt_init();
    memcpy(aes_key,AES_KEY_FIXED,strlen(AES_KEY_FIXED));
    crypt_client();

    return 0;
}

int crypt_client(void)
{
    int ret = -1;
    struct sockaddr_in server_addr;
    unsigned char buf[sizeof(struct in_addr)];
    uint8_t *recv_buf;
    ssize_t len;
    struct pollfd fds[2];
    char input_buf[12];

    printf("run client:\r\n");

    recv_buf = (uint8_t *)calloc(1,2048);
    if(!recv_buf)
    {
        printf("calloc failed..\r\n");
        return ret;
    }

    sock_fd = socket(AF_INET,SOCK_STREAM,0);
    if(sock_fd < 0)
    {
        perror("socket ");
        return ret;
    }

    ret = inet_pton(AF_INET,SERVER_ADDR,buf);
    if(ret != 1)
    {
        printf("inet_pton error..\r\n");
        goto close_sock_fd;
    }
    server_addr.sin_addr.s_addr = ((struct in_addr *)buf)->s_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(SERVER_PORT);
    ret = connect(sock_fd,(const struct sockaddr *)&server_addr,sizeof(server_addr));
    if(ret < 0)
    {
        perror("connect ");
        goto close_sock_fd;
    }

    fds[0].fd = 0;
    fds[0].events = POLLIN;
    fds[1].fd = sock_fd;
    fds[1].events = POLLIN;
    for(;;)
    {
        ret = poll(fds,sizeof(fds) / sizeof(*fds),-1);
        if(ret < 0)
        {
            perror("poll ");
            goto close_sock_fd;
        }
        else if(!ret)
        {
            /*timeout*/
        }
        else
        {
            for(int i = 0;i < sizeof(fds) / sizeof(*fds);i++)
            {
                if(fds[i].revents & fds[i].events)
                {
                    if(fds[i].fd == 0)
                    {
                        len = read(0,input_buf,sizeof(input_buf));
                        if(len < 0)
                        {
                            if((len != EINTR) && (len != EAGAIN) && (len != EWOULDBLOCK))
                            {
                                perror("read ");
                                goto close_sock_fd;
                            }
                        }
                        else if(!len)
                        {

                        }
                        else
                        {
                            unsigned char *ciphertext;
                            ciphertext = (unsigned char *)calloc(1,len);
                            aes_encrypt((const unsigned char *)input_buf,len,aes_key,ciphertext);
                            
                            send(sock_fd,ciphertext,len,0);
                            free(ciphertext);
                        }
                    }
                    else if(fds[i].fd == sock_fd)
                    {
                        memset(recv_buf,0,RECV_BUF_SIZE);
                        len = recv(sock_fd,recv_buf,RECV_BUF_SIZE,0);
                        if(len < 0)
                        {
                            if((len != EINTR) && (len != EAGAIN) && (len != EWOULDBLOCK))
                            {
                                perror("recv ");
                                goto close_sock_fd;
                            }
                            continue;
                        }
                        else if(!len)
                        {
                            printf("server shutdown..\r\n");
                            goto close_sock_fd;
                        }
                        else
                        {
                            unsigned char *plaintext;
                            plaintext = (unsigned char *)calloc(1,len + 1);
                            aes_decrypt(recv_buf,len,aes_key,plaintext);

                            plaintext[len] = '\0';
                            printf("recv:%s\r\n",plaintext);
                            frame_parse(plaintext,len);
                            free(plaintext);
                        }
                    }
                }
            }
        }
    }

    ret = 0;
close_sock_fd:
    close(sock_fd);

    return ret;
}

int frame_parse(const unsigned char *cmd,int len)
{
    char *cmd_s;
    char *ack;
    unsigned char *ciphertext;
    int plaintext_len;
    unsigned char *plaintext;

    cmd_s = "rsa key";
    if(!strncmp((const char *)cmd,cmd_s,strlen(cmd_s)))
    {
        rsa_key = rsa_get_public_from_mem(cmd + strlen(cmd_s),len - strlen(cmd_s));
        if(rsa_key)
        {
            printf("get rsa public key ok..\r\n");

            ack = "public key ok";
            ciphertext = (unsigned char *)calloc(1,strlen(ack));
            aes_encrypt((const unsigned char *)ack,strlen(ack),aes_key,ciphertext);
            send(sock_fd,ciphertext,strlen(ack),0);
            free(ciphertext);
        }
    }
    cmd_s = "aes key";
    if(!strncmp((const char *)cmd,cmd_s,strlen(cmd_s)))
    {
        plaintext = (unsigned char *)calloc(1,AES_KEY_LEN);
        plaintext_len = rsa_public_decrypt(cmd + strlen(cmd_s),rsa_key,plaintext);
        if(plaintext_len < 0)
        {
            printf("rsa_public_decrypt failed..\r\n");
            return -1;
        }
        printf("plaintext_len:%d\r\n",plaintext_len);
        printf("plaintext:\r\n");
        for(int i = 0;i < AES_KEY_LEN;i++)
        {
            if(i && (!(i % 10)))
            {
                printf("\r\n");
            }
            printf("%02x ",plaintext[i]);
        }
        printf("\r\n");

        //验签
        if(1 != rsa_verify(rsa_key,plaintext,plaintext_len,(unsigned char *)(cmd + strlen(cmd_s) + CIPHERTEXT_LEN(rsa_key)),CIPHERTEXT_LEN(rsa_key)))
        {
            printf("rsa_verify failed..\r\n");
        }
        else
        {
            printf("rsa_verify ok..\r\n");
        }

        ack = "aes key ok";
        ciphertext = (unsigned char *)calloc(1,strlen(ack));
        aes_encrypt((const unsigned char *)ack,strlen(ack),aes_key,ciphertext);
        send(sock_fd,ciphertext,strlen(ack),0);
        free(ciphertext);

        //更换密钥
        memcpy(aes_key,plaintext,AES_KEY_LEN);
    }

    return 0;
}
